import numpy as np
from ikpy.chain import Chain
from ikpy.link import DHLink
from ikpy.urdf.URDF import get_urdf_parameters

from ezgl import Matrix4x4


class IKSolver:

    def __init__(self, urdf_path=None, base_elements: list=None, last_joint:str=None): 
        if urdf_path is None:            
            self.robot = Chain(name='ur5_dh', links=[
                DHLink(name="joint1",   d=0.089159, a=0,         alpha=np.pi/2, theta=0),
                DHLink(name="joint2",   d=0,        a=-0.425,    alpha=0,       theta=0),
                DHLink(name="joint3",   d=0,        a=-0.39225,  alpha=0,       theta=0),
                DHLink(name="joint4",   d=0.10915,  a=0,         alpha=np.pi/2, theta=0),
                DHLink(name="joint5",   d=0.09465,  a=0,         alpha=-np.pi/2,theta=0),
                DHLink(name="joint6",   d=0.0823,   a=0,         alpha=0,       theta=0),
            ])
            self._active_links_mask = np.array([1, 1, 1, 1, 1, 1], dtype=bool)
        else:
            links = get_urdf_parameters(urdf_file=urdf_path, base_elements=base_elements)
            for i in range(len(links)):
                if links[i].name == last_joint:
                    break
            links = links[:i+1]  # 截断到指定的关节
            
            self._active_links_mask = np.array([link.joint_type != "fixed" for link in links], dtype=bool)
            self.robot = Chain(name='robot', links=links, active_links_mask=self._active_links_mask)
    
    def _full_angles(self, angles):
        if len(angles) == len(self._active_links_mask):
            return angles
        else:
            full_angles = np.zeros(len(self._active_links_mask))
            full_angles[self._active_links_mask] = angles
            return full_angles

    def _active_angles(self, angles):
        if len(angles) == len(self._active_links_mask):
            return angles[self._active_links_mask]
        else:
            return angles
    
    def forward_kinematics(self, angles):
        return self.robot.forward_kinematics(self._full_angles(angles)) # 4*4 matrix
    
    def inverse_kinematics(self, pose: Matrix4x4, initial_guess):
        try:
            x = self.robot.inverse_kinematics_frame(target=pose.matrix44, orientation_mode="all", initial_position=self._full_angles(initial_guess))
            x = self._active_angles(x)  # 返回活动关节的角度
        except Exception as e:
            print(f"Inverse kinematics failed: {e}")
            x = None
        return x


if __name__ == '__main__':
    pose = Matrix4x4.fromTranslation(0.5, 0.2, 0.4)
    
    solver = IKSolver(urdf_path=r"H:\myapps\ezrobot\ezrobot\assets\ur5\ur5.urdf", base_elements=["world"], last_joint="ee_link")
    print(solver._active_links_mask)
    j = solver.inverse_kinematics(pose, [0, 0, 0, 0, 0, 0])  # 初始猜测为全零
    from ezrobot.sim import RobotClient
    client = RobotClient()
    client.set_joint(j)